

from utils.data import ListDataset 
from torch.utils.data import DataLoader
import torch 
import os 
import matplotlib.pyplot as plt 
from utils.model import YoloModel
from utils.loss import compute_loss
def main():

    dataset = ListDataset("data")
    dataloader = DataLoader(dataset, 32, shuffle=True, num_workers=8, collate_fn=dataset.collate_fn)
    device = torch.device("cpu")
    model = YoloModel() 
    model.train() 
    model.to(device)
    if os.path.exists("ckpt/model.pt"):
        model.load_state_dict(torch.load("ckpt/model.pt"))
    for key, var in model.named_parameters():
        if var.dtype != torch.float32:continue # BN统计计数无梯度
        if ".out." in key: # 仅有最后一层有out
            var.requires_grad = True
        else:
            var.requires_grad = False
    #optim = torch.optim.Adam(model.parameters(), 1e-3, weight_decay=1e-3)
    optim = torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()))
    n = 0
    print("样本数量", len(dataset))
    for e in range(30):
        for temp, imgs, targets in dataloader:
            imgs = imgs.to(device) 
            targets = targets.to(device)
            outs = model(imgs) 
            loss = compute_loss(outs, targets) 
            loss.backward() 
            optim.step()
            optim.zero_grad() 
            if n % 50 == 0:
                print(e, n, loss)
                torch.save(model.state_dict(), f"ckpt/model.pt")
            n += 1
        torch.save(model.state_dict(), f"ckpt/{e}.pt")
        print(loss)
#nohup /home/yuzy/software/anaconda39/bin/python yolo.train.py > ckpt/yolo.log 2>&1 &
if __name__ == "__main__":
    main()
